# -*- coding: utf-8 -*-
"""
@Time    : 2024/10/22 11:44 
@Author  : ZhangShenao 
@File    : initiator.py 
@Desc    : 初始化器

执行BaseChatModel、Embeddings、VectorStore等组件的初始化工作
"""
import dotenv
from FlagEmbedding import FlagReranker
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseChatModel
from langchain_core.vectorstores import VectorStore
from langchain_pinecone import PineconeVectorStore


def init_components() -> (BaseChatModel, Embeddings, VectorStore, FlagReranker):
    """
    初始化BaseChatModel、Embeddings、VectorStore、Reranker等组件
    """

    print("组件初始化开始执行!")

    # 加载环境变量
    dotenv.load_dotenv()

    # 创建ChatZhipuAI
    llm = ChatZhipuAI(model='glm-4-flash')

    # 创建ZhipuAIEmbeddings,指定dimensions为1024
    embeddings = ZhipuAIEmbeddings(model="embedding-3", dimensions=1024)

    # 创建PineConeVectorStore
    vector_store = PineconeVectorStore(
        index_name='rag-app',  # 指定索引名称
        embedding=embeddings,  # 指定Embedding模型
        namespace='llm-ops',  # 指定Namespace
        text_key='text',  # 指定文本属性key
    )

    # 创建FlagReranker
    reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True)

    # 返回组件
    print("组件初始化执行完成!")
    return llm, embeddings, vector_store, reranker
